% @author:夜剑听雨
% E-mail:2530595378@qq.com
% 2025.08.29
% 频散点拾取程序
clear; close all;

% 设置文件夹路径
dataFolder = './data/data_mat';
outputFolder = './save_dir/pick';
% 创建输出文件夹
if ~exist(outputFolder, 'dir')
    mkdir(outputFolder);
end
%% 设置参数
%  务必和function MouseDownCallback中的数值保持一致，需要手动修改
RangeFreq = [5,65];
ParaVelocity = [50 800 1];
nu = 0;
dataMode = 0;

% 明确拾取范围变量
fMin = RangeFreq(1);
fMax = RangeFreq(2);
vMin = ParaVelocity(1);
vMax = ParaVelocity(2);

% 获取所有 .mat 文件
fileNames = dir(fullfile(dataFolder, '*.mat'));
fileNames = {fileNames.name};
%% 主处理循环
for i = 1:length(fileNames)
    % 加载数据
    filePath = fullfile(dataFolder, fileNames{i});
    if ~exist(filePath, 'file')
        fprintf('文件不存在: %s\n', filePath);
        continue;
    end
    load(filePath);

    % 数据准备
    data = shot.ZZ;
    offset = shot.offset;
    dx = shot.dx;
    dt = shot.dt;

    %% 频散变换
    fs = 1/dt;
    traceNum = size(data,2);
    r = offset + (0:traceNum-1)*dx;


    % 计算频散能量
    [E, freq, v] = phaseShiftTransform(data,fs,r,RangeFreq,ParaVelocity,nu,dataMode);

    %% 绘制初始频散图
    fh = figure('Name',['Processing: ' fileNames{i}],'NumberTitle','off');
    ax = gca;
    [X,Y] = meshgrid(freq,v);
    [~,ch] = contourf(X,Y,E,100);
    set(ch, 'LineStyle', 'none');
    set(ch, 'EdgeColor', 'none');
    colormap(jet);
    shading flat;
    colorbar;
    set(gca,'YDir','normal', 'LineWidth',1, 'FontSize',12);
    xlim([fMin fMax]);
    xticks(fMin:10:fMax);
    ylim([vMin vMax]);
    yticks(vMin:50:vMax);
    title('Phase Shift Transform','FontSize',12);
    xlabel('Frequency(Hz)','FontSize',12);
    ylabel('Phase velocity(m/s)','FontSize',12);
    set(gca,'LooseInset',get(gca,'TightInset')+0.01);

    % 强制设置坐标轴范围确保一致性
    axis([fMin fMax vMin vMax]);
   
    %% 初始化存储结构
    Freq_length = fMax - fMin + 1;
    picked_data = struct(...
        'original_filename', fileNames{i},...
        'frequency', (fMin:1:fMax)',...
        'velocity1', nan(Freq_length,1),...
        'velocity2', nan(Freq_length,1),...
        'velocity3', nan(Freq_length,1));

    %% 频散曲线拾取主循环
    for curveIdx = 1:3
        % 初始化当前曲线参数
        setappdata(fh, 'currentCurve', curveIdx);
        setappdata(fh, 'x_pick', []);
        setappdata(fh, 'y_pick', []);
        setappdata(fh, 'KeyPressed', false);

        % 设置回调函数
        set(fh, 'KeyPressFcn', @KeyPressCallback);
        set(fh, 'WindowButtonDownFcn', @MouseDownCallback);

        % 显示提示信息
        msg = ['请拾取第' num2str(curveIdx) '条频散曲线（按空格结束）'];
        uicontrol('Style','text',...
            'String',msg,...
            'Units','normalized',...
            'Position',[0.3 0.95 0.4 0.04],...
            'FontSize',10,...
            'BackgroundColor',get(fh,'Color'));

        % 等待用户完成拾取
        uiwait(fh);

        % 获取拾取点
        x_pick = getappdata(fh, 'x_pick');
        y_pick = getappdata(fh, 'y_pick');

        % 数据处理
        if numel(x_pick) < 2
            warning('曲线%d: 有效点不足，已跳过', curveIdx);
            continue;
        end

        % 限制频率范围
        validIdx = (x_pick >= fMin) & (x_pick <= fMax);
        x_pick = x_pick(validIdx);
        y_pick = y_pick(validIdx);

        % 频率对齐到1Hz间隔
        x_pick_rounded = round(x_pick);

        % 去除重复点
        [~, uniqueIdx] = unique(x_pick_rounded);
        x_pick = x_pick_rounded(uniqueIdx);
        y_pick = y_pick(uniqueIdx);

        % 排序
        [x_pick_sorted, sort_idx] = sort(x_pick);
        y_pick_sorted = y_pick(sort_idx);

        % 截断频率范围
        f_min_rounded = floor(x_pick_sorted(1));
        f_max_rounded = ceil(x_pick_sorted(end));

        % 确保在有效范围内
        f_min_rounded = max(f_min_rounded, fMin);
        f_max_rounded = min(f_max_rounded, fMax);

        % 创建目标频率数组
        freq_target = f_min_rounded:1:f_max_rounded;

        % 插值处理
        if numel(x_pick) > 1
            vel_interp = interp1(x_pick_sorted, y_pick_sorted, freq_target, 'pchip', 'extrap');
            vel_interp = min(max(vel_interp, vMin), vMax);  % 速度限幅
        else
            vel_interp = nan(size(freq_target));
        end

        % 保存数据
        switch curveIdx
            case 1
                picked_data.velocity1 = zeros(length(picked_data.frequency), 1);
                idx = find(picked_data.frequency >= f_min_rounded & picked_data.frequency <= f_max_rounded);
                picked_data.velocity1(idx) = vel_interp;
            case 2
                picked_data.velocity2 = zeros(length(picked_data.frequency), 1);
                idx = find(picked_data.frequency >= f_min_rounded & picked_data.frequency <= f_max_rounded);
                picked_data.velocity2(idx) = vel_interp;
            case 3
                picked_data.velocity3 = zeros(length(picked_data.frequency), 1);
                idx = find(picked_data.frequency >= f_min_rounded & picked_data.frequency <= f_max_rounded);
                picked_data.velocity3(idx) = vel_interp;
        end

        % 绘制结果
        hold(ax, 'on');
        plot(ax, freq_target, vel_interp, 'k.', 'MarkerSize', 4); % 使用黑色散点图
        hold(ax, 'off');
        drawnow;

        % 询问是否继续
        if curveIdx < 3
            choice = questdlg(sprintf('是否继续拾取第%d条曲线？', curveIdx+1),...
                '确认', '继续', '跳过', '继续');
            if strcmp(choice, '跳过')
                break;
            end
        end
    end

    %% 最终处理
    % 保存数据
    [~, baseName] = fileparts(fileNames{i});
    save(fullfile(outputFolder, [baseName '_picked.mat']), 'picked_data');

    % 保存最终图像
    finalFigPath = fullfile(outputFolder, strrep(fileNames{i}, '.mat', '_final.jpg'));
    saveas(fh, finalFigPath, 'jpg');
    close(fh);
end


function KeyPressCallback(src, event)
fh = gcbf;
if strcmp(event.Key, 'space')
    setappdata(fh, 'KeyPressed', true);
    uiresume(fh);

    % 删除临时标记
    ax = get(fh, 'CurrentAxes');
    curveIdx = getappdata(fh, 'currentCurve');
    delete(findobj(ax, 'Tag', ['pickMark' num2str(curveIdx)]));

    % 弹出确认对话框
    response = questdlg('是否结束本次拾取？', '确认', '是', '否', '否');
    if strcmp(response, '否')
        setappdata(fh, 'KeyPressed', false);
        % 恢复拾取状态
        curveIdx = getappdata(fh, 'currentCurve');
        msg = ['请继续拾取第' num2str(curveIdx) '条频散曲线（按空格结束）'];
        uicontrol('Style','text',...
            'String',msg,...
            'Units','normalized',...
            'Position',[0.3 0.95 0.4 0.04],...
            'FontSize',10,...
            'BackgroundColor',get(fh,'Color'));
        uiwait(fh);
    end
end
end

function MouseDownCallback(src, event)
fh = gcbf;
curveIdx = getappdata(fh, 'currentCurve');

if strcmp(get(src, 'SelectionType'), 'normal')
    ax = get(fh, 'CurrentAxes');
    cp = get(ax, 'CurrentPoint');
    x = cp(1,1);
    y = cp(1,2);

    % 明确拾取范围变量（需要在函数中重新定义）
    RangeFreq = [5,65];
    ParaVelocity = [50 800 1];
    fMin = RangeFreq(1);
    fMax = RangeFreq(2);
    vMin = ParaVelocity(1);
    vMax = ParaVelocity(2);

    % 有效性检查
    if x < fMin || x > fMax || y < vMin || y > vMax
        return;
    end

    % 调试输出
    fprintf('Clicked at: (%.2f, %.2f)\n', x, y);

    % 存储数据
    x_pick = getappdata(fh, 'x_pick');
    y_pick = getappdata(fh, 'y_pick');
    x_pick(end+1) = x;
    y_pick(end+1) = y;
    setappdata(fh, 'x_pick', x_pick);
    setappdata(fh, 'y_pick', y_pick);

    % 绘制标记
    hold(ax, 'on');
    ph = plot(ax, x, y, 'ko',...
        'MarkerSize', 5,...
        'MarkerFaceColor', 'k',...
        'Tag', ['pickMark' num2str(curveIdx)]);

    hold(ax, 'off');
    drawnow;
end
end

